import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt

# 分割数据为训练集和测试集
from sklearn.model_selection import train_test_split

# 读取输入数据
data_df = pd.read_csv("./data/source.csv")
inputs = data_df[["集水面积", "Sr", "Ks", "蒸发量", "降雨量(24h)"]].values
true_outputs = data_df["洪峰流量"].values


X_train, X_test, y_train, y_test = train_test_split(
    inputs, true_outputs, test_size=0.2, random_state=42
)

# 初始化随机森林模型
rf_model = RandomForestRegressor(n_estimators=100, random_state=42)

# 训练模型
rf_model.fit(X_train, y_train)

# 预测训练集和测试集
train_predictions = rf_model.predict(X_train)
test_predictions = rf_model.predict(X_test)

# 计算损失
train_mse = mean_squared_error(y_train, train_predictions)
test_mse = mean_squared_error(y_test, test_predictions)

print(f"Training MSE: {train_mse:.4f}")
print(f"Testing MSE: {test_mse:.4f}")

# 提取随机森林中的参数（这里近似为特征重要性的加权）
parameters = rf_model.feature_importances_

# 给33个参数指定名称
parameter_names = [
    "B",
    "alpha0",
    "JL",
    "KKS",
    "KKSS",
    "CS",
    "z1",
    "z2",
    "z3",
    "ni1",
    "ni2",
    "ni3",
    "H01",
    "H02",
    "H03",
    "Hc1",
    "Hc2",
    "Hc3",
    "Hr1",
    "Hr2",
    "Hr3",
    "delta1",
    "delta2",
    "delta3",
    "LS",
    "LSS",
    "L",
    "Gam1",
    "Gam2",
    "Gam3",
    "EC",
    "KKG",
    "LG",
]

# 如果参数数量不足33个，则扩展为默认值
if len(parameters) < 33:
    additional_params = [0] * (33 - len(parameters))
    parameters = np.concatenate([parameters, additional_params])

# 保存为DataFrame
parameter_df = pd.DataFrame(
    {
        "Parameter": parameter_names,
        "Value": parameters[:33],  # 确保只保存33个参数
    }
)

# 保存为CSV文件
output_filename = "./output/parameters.csv"
parameter_df.to_csv(output_filename, index=False)

print(f"Parameters saved to {output_filename}")

# 可视化参数值
plt.figure(figsize=(12, 8))
plt.barh(parameter_df["Parameter"], parameter_df["Value"], color="skyblue")
plt.title("Parameter Values")
plt.xlabel("Value")
plt.ylabel("Parameter")
plt.grid(True)
plt.tight_layout()
plt.show()
